package com.yujiahao;

import org.apache.flink.kinesis.shaded.com.amazonaws.ClientConfiguration;
import org.apache.flink.kinesis.shaded.com.amazonaws.auth.DefaultAWSCredentialsProviderChain;
import org.apache.flink.kinesis.shaded.com.amazonaws.services.kinesis.AmazonKinesis;
import org.apache.flink.kinesis.shaded.com.amazonaws.services.kinesis.AmazonKinesisClientBuilder;
import org.apache.flink.kinesis.shaded.com.amazonaws.services.kinesis.model.*;
import org.apache.flink.kinesis.shaded.software.amazon.awssdk.services.kinesis.model.DescribeStreamResponse;
import org.junit.Before;
import org.junit.Test;

import java.math.BigInteger;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.stream.Collectors;

import static org.junit.Assert.assertEquals;

/**
 * @Description:kinesis 生产者
 * @Author: yujiahao
 * @Date: 2023/7/31 14:52
 * @Version: 1.0
 */

public class KinesisProducer {


    private static AmazonKinesis client;
    private String streamName = "test-kds-api";
    private String sequenceNumberOfPreviousRecord;
    DescribeStreamRequest describeStreamRequest;

    /**
     * 创建客户端对象
     */
    @Before
    public void createClient() {
        //构建客户端对象
        AmazonKinesisClientBuilder clientBuilder = AmazonKinesisClientBuilder.standard();
        ClientConfiguration clientConfiguration = new ClientConfiguration();
        clientBuilder.setCredentials(new DefaultAWSCredentialsProviderChain());
        clientBuilder.setClientConfiguration(clientConfiguration);
        client = clientBuilder.build();
        describeStreamRequest = new DescribeStreamRequest().withStreamName(streamName);



    }

    /**
     * 为每个请求发送单个记录
     * 以下代码创建十个数据记录，分布在两个分区键上，并将它们放入名为 的流中myStreamName。
     */
    @Test
    public void putRecord() {

        for (int j = 0; j < 10; j++) {
            // 创建 PutRecordRequest 请求
            PutRecordRequest putRecordRequest = new PutRecordRequest();

            // 设置要写入的 Kinesis 流名称
            putRecordRequest.setStreamName(streamName);

            // 设置要写入的数据。这里使用 String.format 将数据格式化为 "testData-j" 的形式，并将其转换为 ByteBuffer。
            putRecordRequest.setData(ByteBuffer.wrap(String.format("testData-%d", j).getBytes()));

            // 设置数据的分区键（Partition Key）。这里使用 String.format 将分区键格式化为 "partitionKey-j/5" 的形式。
            //  j/5 的结果是用来做分区的，这里有2个分区 0 和 1
            /*
            对于 j = 0，分区键为 "partitionKey-0"
            对于 j = 1，分区键为 "partitionKey-0"
            对于 j = 2，分区键为 "partitionKey-0"
            对于 j = 3，分区键为 "partitionKey-0"
            对于 j = 4，分区键为 "partitionKey-0"
            对于 j = 5，分区键为 "partitionKey-1"
            对于 j = 6，分区键为 "partitionKey-1"
            对于 j = 7，分区键为 "partitionKey-1"
            对于 j = 8，分区键为 "partitionKey-1"
            对于 j = 9，分区键为 "partitionKey-1"
             */
            //Kinesis分区键是必须要有的，不然数据懵了不知道自己要去哪，这里就是自己简单设置了一下
            //MD5 计算该记录的哈希键值做分区键
            putRecordRequest.setPartitionKey(String.format("partitionKey-%d", j / 5));

            // 设置顺序号，用于记录排序。这里使用之前写入的记录的顺序号作为当前记录的排序依据。（分区内，直接不能设置全局顺序）
            putRecordRequest.setSequenceNumberForOrdering(sequenceNumberOfPreviousRecord);

            // 使用 Amazon Kinesis 客户端的 putRecord() 方法将记录写入流中，并返回 PutRecordResult 结果对象。
            PutRecordResult putRecordResult = client.putRecord(putRecordRequest);

            // 从 PutRecordResult 中获取写入记录后的顺序号，用于下一次写入记录时设置顺序号。这是获取分区内顺序号
            sequenceNumberOfPreviousRecord = putRecordResult.getSequenceNumber();
            //获取并打印一下分区号看看
            String partitionKey = putRecordRequest.getPartitionKey();
            System.out.println("Data: " + putRecordRequest.getData() + ", Partition Key: " + partitionKey);
            //分片id打印一下
            System.out.println("分片的 ID: " + putRecordResult.getShardId());

        }
    }


    /**
     * 使用PutRecords批处理，这个用的多
     * 以下代码使用顺序分区键创建 100 个数据记录，并将它们放入名为 的流中DataStream。
     * PutRecords包括一个响应数组 Records。响应数组中的每个记录使用自然排序（从请求和响应的顶部到底部）直接与请求数组中的记录相关联。
     * 响应Records数组始终包含与请求数组相同数量的记录。
     */
    @Test
    public void putRecords() {

        // 创建批量写入请求对象 PutRecordsRequest
        PutRecordsRequest putRecordsRequest = new PutRecordsRequest();
        // 设置要写入的目标流的名称
        putRecordsRequest.setStreamName(streamName);
        // 创建一个列表用于存储多个数据记录
        List<PutRecordsRequestEntry> putRecordsRequestEntryList = new ArrayList<>();
        for (int i = 0; i < 100; i++) {
            // 创建数据记录请求对象 PutRecordsRequestEntry
            PutRecordsRequestEntry putRecordsRequestEntry = new PutRecordsRequestEntry();
            //设置和单条写入都一样没啥滑头
            putRecordsRequestEntry.setData(ByteBuffer.wrap(String.valueOf(i).getBytes()));
            putRecordsRequestEntry.setPartitionKey(String.format("partitionKey-%d", i));
            putRecordsRequestEntryList.add(putRecordsRequestEntry);
        }
        // 将包含多条数据记录的列表设置到批量写入请求对象中
        putRecordsRequest.setRecords(putRecordsRequestEntryList);
        // 使用 putRecords() 方法批量写入数据记录，并获取写入结果
        PutRecordsResult putRecordsResult = client.putRecords(putRecordsRequest);
        //打印看看
        System.out.println("Put Result" + putRecordsResult);

    /**
    todo 注意默认情况下，请求中单个记录的失败不会停止请求中后续记录的处理PutRecords。 这意味着响应Records数组包括成功处理和不成功处理的记录。必须检测未成功处理的记录并将其包含在后续调用中
        成功的记录包括SequenceNumber和 ShardID值，不成功的记录包括 ErrorCode和ErrorMessage值。
        该 ErrorCode参数反映错误的类型，可以是以下值之一：ProvisionedThroughputExceededException或 InternalFailure。
        ErrorMessage提供有关异常的更多详细信息， ProvisionedThroughputExceededException包括帐户 ID、流名称和受到限制的记录的分片 ID。

   todo  ① PutRecords 请求语法示例

        {
            "Records": [
                {
                "Data": "XzxkYXRhPl8w",
                "PartitionKey": "partitionKey1"
                },
                {
                "Data": "AbceddeRFfg12asd",
                "PartitionKey": "partitionKey1"
                },
                {
                "Data": "KFpcd98*7nd1",
                "PartitionKey": "partitionKey3"
                }
            ],
            "StreamName": "myStream"
        }

    todo ② PutRecords 响应语法示例 批处理有错误

        {
            "FailedRecordCount”: 1,
            "Records": [
                {
                "SequenceNumber": "21269319989900637946712965403778482371",
                "ShardId": "shardId-000000000001"

                },
                {
                “ErrorCode":”ProvisionedThroughputExceededException”,
                “ErrorMessage": "Rate exceeded for shard shardId-000000000001 in stream exampleStreamName under account 111111111111."

                },
                {
                "SequenceNumber": "21269319989999637946712965403778482985",
                "ShardId": "shardId-000000000002"
                }
            ]
        }

 todo 未成功处理的记录可以包含在后续 PutRecords请求中。首先检查 FailedRecordCount其中的参数， putRecordsResult确认请求中是否有失败记录。
  如果是这样，则每个putRecordsEntry有一个 ErrorCode不是的null都应该添加到后续请求中。有关此类处理程序的示例，请参阅以下代码。
     */
        while (putRecordsResult.getFailedRecordCount() > 0) {
            // 创建一个列表用于存储写入失败的数据记录
            final List<PutRecordsRequestEntry> failedRecordsList = new ArrayList<>();
            // 获取写入结果中的数据记录列表
            final List<PutRecordsResultEntry> putRecordsResultEntryList = putRecordsResult.getRecords();

            // 遍历写入结果列表，检查每条数据记录是否写入失败
            for (int i = 0; i < putRecordsResultEntryList.size(); i++) {
                // 获取原始的数据记录请求对象
                final PutRecordsRequestEntry putRecordRequestEntry = putRecordsRequestEntryList.get(i);
                // 获取写入结果中的数据记录对象
                final PutRecordsResultEntry putRecordsResultEntry = putRecordsResultEntryList.get(i);

                //从结果获取写失败记录，（如果存在错误代码就是写失败了），然后将数据添加到失败记录中
                if (putRecordsResultEntry.getErrorCode() != null) {
                    failedRecordsList.add(putRecordRequestEntry);
                }
            }
            // 将失败列表赋值给数据记录列表，用于下一次重新写入
            putRecordsRequestEntryList = failedRecordsList;
            // 将新的数据记录列表设置到批量写入请求对象中
            putRecordsRequest.setRecords(putRecordsRequestEntryList);
            // 使用 putRecords() 方法重新批量写入数据记录，并获取写入结果
            putRecordsResult = client.putRecords(putRecordsRequest);
        }


    }


}
